Add torch_compile flag for training networks#28
Conversation
|
@juliusberner Could you please take a look at my PR? Thanks! |
wenxin0319
force-pushed
the
main
branch
3 times, most recently
from
eb763ac to
78e9f0d
Compare
juliusberner
left a comment
juliusberner
left a comment
There was a problem hiding this comment.
Thanks for the edits! I posted a couple of follow-up comments and also kicked off the CI!
Also, note that per our contribution guidelines "Commits must have verified signatures.".
| if hasattr(self.net, "init_preprocessors") and self.config.enable_preprocessors: | ||
| self.net.init_preprocessors() | ||
|
|
||
| # Preprocessor sub-objects of ``net`` to also compile (see compile_dict). |
There was a problem hiding this comment.
Can we define this at the top and also use _PREPROCESSOR_ATTRS in on_train_begin, so that this tuple will be the only source-of-truth for the preproc. attributes?
| _PREPROCESSOR_ATTRS = ("vae", "text_encoder", "image_encoder") | ||
|
|
||
| @property | ||
| def compile_dict(self) -> dict: |
There was a problem hiding this comment.
Do you think we could take the model_dict as base dict (which already contains, e.g., fake score and discriminator for DMD2) and
- remove the EMA from the
ema_dict? - add the teacher to the base compile dict, guarded by an
getattr(self, "teacher") is not None(similar to thefsdp_dict) - add the preprocessors as done below
| return compile_dict | ||
|
|
||
| @staticmethod | ||
| def _object_compile_targets(name: str, obj) -> dict: |
There was a problem hiding this comment.
Can we just inline this in the compile_dict method, since it's not used anywhere else?
wenxin0319
force-pushed
the
main
branch
4 times, most recently
from
5100a24 to
b7226d8
Compare
|
@juliusberner Thank you for reviewing it; in the latest commit, i am trying to make the code as concise as i can. let me know if you have any comments. |
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@wenxin0319 Thanks a lot, I've just started the final coderabbit review.
After that, we can merge :) |
WalkthroughThis PR introduces optional torch.compile support for FastGen models. A new configuration field enables torch compilation, a central preprocessor attribute list unifies device placement logic, and a new apply_torch_compile() method compiles trainable modules after distributed wrapping. The trainer invokes compilation post-DDP/FSDP, and comprehensive tests validate behavior across SFT and DMD2 models including edge cases and training integration. Changestorch.compile Feature
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/test_torch_compile.py`:
- Line 174: The test unpacks two values from model.single_train_step into
loss_map and outputs but never uses outputs; to silence RUF059, change the
unpacking in tests/test_torch_compile.py where single_train_step is called
(references: model.single_train_step) to either discard the second value by
assigning it to _ (e.g., loss_map, _ = ...) or only capture the first element
(e.g., loss_map = model.single_train_step(...)[0]); apply the same change for
both occurrences mentioned (the calls at the two test locations).
- Around line 120-124: The test_dmd2_compile_disabled test is missing an
assertion for the discriminator; add an assertion that _is_compiled returns
False for dmd2_model_not_compiled.discriminator (mirroring the enabled-path test
coverage), i.e., extend test_dmd2_compile_disabled to include assert not
_is_compiled(dmd2_model_not_compiled.discriminator) so the disabled-path
validates the same component as the enabled-path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Enterprise
Run ID: 41b1322a-6dbf-4043-809f-ca12de965606
📒 Files selected for processing (4)
fastgen/configs/config.pyfastgen/methods/model.pyfastgen/trainer.pytests/test_torch_compile.py
Add an optional torch_compile_mode config flag ("default",
"reduce-overhead", "max-autotune"; None disables). Networks to compile
are collected in a compile_dict (derived from model_dict minus EMA, plus
the teacher and the net's preprocessors) and compiled in place via
nn.Module.compile by the trainer, after DDP/FSDP wrapping.
hi, I have added it to image/video inference.py; and the commit is verified right now |
Thanks a lot! Could you address the two CodeRabbit comments above and my remaining comment? |
Greptile SummaryThis PR adds an opt-in
Confidence Score: 4/5The core compile-and-train path is correct and well-tested; the one remaining concern is the EMA network being excluded from compilation even during the inference path where it is the active student network. When
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
subgraph Training["Training Path (Trainer.run)"]
T1[model.on_train_begin] --> T2[DDP / FSDP wrapping]
T2 --> T3[model.apply_torch_compile]
T3 --> T4{torch_compile_mode?}
T4 -- None --> T5[no-op]
T4 -- mode string --> T6[compile model_dict minus EMA]
T6 --> T7[compile teacher if present]
T7 --> T8[compile preprocessor sub-modules]
end
subgraph Inference["Inference Path (setup_inference_modules)"]
I1[cleanup_unused_modules] --> I2[resolve teacher / student refs]
I2 --> I3[init_preprocessors]
I3 --> I4[model.apply_torch_compile]
I4 --> I5{torch_compile_mode?}
I5 -- None --> I6[no-op]
I5 -- mode string --> I7[compile model_dict minus EMA + teacher + preprocessors]
end
subgraph Guard["Idempotency Guard"]
G1{_compiled_call_impl already set?} -- yes --> G2[skip, log warning]
G1 -- no --> G3[module.compile mode=mode]
end
T6 --> G1
T7 --> G1
T8 --> G1
I7 --> G1
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
subgraph Training["Training Path (Trainer.run)"]
T1[model.on_train_begin] --> T2[DDP / FSDP wrapping]
T2 --> T3[model.apply_torch_compile]
T3 --> T4{torch_compile_mode?}
T4 -- None --> T5[no-op]
T4 -- mode string --> T6[compile model_dict minus EMA]
T6 --> T7[compile teacher if present]
T7 --> T8[compile preprocessor sub-modules]
end
subgraph Inference["Inference Path (setup_inference_modules)"]
I1[cleanup_unused_modules] --> I2[resolve teacher / student refs]
I2 --> I3[init_preprocessors]
I3 --> I4[model.apply_torch_compile]
I4 --> I5{torch_compile_mode?}
I5 -- None --> I6[no-op]
I5 -- mode string --> I7[compile model_dict minus EMA + teacher + preprocessors]
end
subgraph Guard["Idempotency Guard"]
G1{_compiled_call_impl already set?} -- yes --> G2[skip, log warning]
G1 -- no --> G3[module.compile mode=mode]
end
T6 --> G1
T7 --> G1
T8 --> G1
I7 --> G1
Reviews (4): Last reviewed commit: "Fix test fixture ordering to match train..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # Compilation is applied by the trainer after DDP/FSDP wrapping; emulate that here. | ||
| model.apply_torch_compile() |
There was a problem hiding this comment.
The fixtures call model.apply_torch_compile() before model.on_train_begin(), but in production the trainer runs the opposite order: on_train_begin() at line 97 of trainer.py, then apply_torch_compile() at line 124. Because on_train_begin() is what calls init_preprocessors() (creating net.vae, net.text_encoder, etc.) and moves parameters to the target device/dtype, calling apply_torch_compile() first means preprocessors don't exist yet and are never compiled — even when torch_compile_mode is set. Any model that relies on compiled preprocessors would appear to work in these tests while silently skipping preprocessor compilation in production. The same wrong order appears in the dmd2_model_compiled, sft_model_not_compiled, and dmd2_model_not_compiled fixtures.
| Returns: | ||
| Tuple of (teacher, student, vae) - any may be None | ||
| """ | ||
| model.apply_torch_compile() # no-op if torch_compile_mode is None |
There was a problem hiding this comment.
apply_torch_compile() is called here before init_preprocessors() is called a few lines later (line 181). Since preprocessors (net.vae, net.text_encoder, net.image_encoder) only exist after init_preprocessors() runs, apply_torch_compile() won't find them and will skip compiling them. In training this is handled correctly — on_train_begin() initialises preprocessors before apply_torch_compile() is invoked by the trainer. For inference with torch_compile_mode set, users expecting compiled preprocessors will silently get uncompiled versions.
| def apply_torch_compile(self): | ||
| """Compile the training networks in place with torch.compile. | ||
|
|
||
| Called by the trainer after DDP/FSDP wrapping (and after the networks | ||
| have been moved to their device in ``on_train_begin``) so torch.compile | ||
| composes with the distributed wrappers. No-op when | ||
| ``config.torch_compile_mode`` is None. | ||
|
|
||
| The modules compiled are those in model_dict (e.g. net, plus | ||
| fake_score/discriminator for DMD2) minus the EMA networks, plus the | ||
| teacher (if any, cf. fsdp_dict) and the net's preprocessors. | ||
| """ | ||
| mode = self.config.torch_compile_mode | ||
| if mode is None: | ||
| return | ||
|
|
||
| # model_dict contains the trainable networks (incl. EMA); EMA networks are | ||
| # weight-averaged copies that aren't run during training, so drop them. | ||
| # None entries arise when cleanup_unused_modules has already been called. | ||
| modules = {name: net for name, net in self.model_dict.items() if name not in self.ema_dict and net is not None} | ||
| # The teacher is not part of model_dict; add it when present (cf. fsdp_dict). | ||
| if getattr(self, "teacher", None) is not None: | ||
| modules["teacher"] = self.teacher | ||
| # Preprocessors (VAE, text/image encoders) are often lightweight wrappers | ||
| # (e.g. WanVideoEncoder, SDVAE) that are not nn.Modules themselves but hold | ||
| # the actual nn.Module under an attribute; compile those submodules too. | ||
| for name in self._PREPROCESSOR_ATTRS: | ||
| obj = getattr(self.net, name, None) | ||
| if obj is None: | ||
| continue | ||
| if isinstance(obj, torch.nn.Module): | ||
| modules[name] = obj | ||
| else: | ||
| for attr, submodule in getattr(obj, "__dict__", {}).items(): | ||
| if isinstance(submodule, torch.nn.Module): | ||
| modules[f"{name}.{attr}"] = submodule | ||
|
|
||
| for name, module in modules.items(): | ||
| logger.info(f"Applying torch.compile (mode={mode}) to {name}") | ||
| module.compile(mode=mode) |
There was a problem hiding this comment.
apply_torch_compile() has no guard against being called more than once on the same model instance. nn.Module.compile() applied a second time wraps the already-compiled _call_impl, compiling a compiled callable. While there is no code path in the current repo that calls this twice in one session, any future callback or utility bridging the training and inference paths could double-compile. A simple check like if getattr(module, "_compiled_call_impl", None) is not None: continue inside the loop would prevent this.
There was a problem hiding this comment.
@juliusberner thank you for pointing it out, i have fixed both comments
|
@wenxin0319 it looks great! Could you fix the CI, such that we can merge it? |
| # weight-averaged copies that aren't run during training, so drop them. | ||
| # None entries arise when cleanup_unused_modules has already been called. | ||
| modules = {name: net for name, net in self.model_dict.items() if name not in self.ema_dict and net is not None} |
There was a problem hiding this comment.
apply_torch_compile always excludes every name in ema_dict, but at inference time setup_inference_modules (line 175 of inference_utils.py) picks the EMA network as the student: student = getattr(model, model.use_ema[0]) if model.use_ema else model.net. That EMA network is the one actually called for student sampling, yet it is never compiled here. model.net gets compiled but is not invoked; the EMA network that is invoked stays uncompiled. Any user who sets torch_compile_mode expecting compiled inference will silently get no speedup when use_ema is non-empty (the default for most production configs).
The property inference_net (line 729) confirms the same selection: it returns self.use_ema[0] when EMA is active. Since the compile exclusion is designed for training semantics (EMA weights are updated by averaging, not by a forward pass), it should not unconditionally apply to the inference path. One option is to additionally compile whichever network inference_net resolves to when called from setup_inference_modules.
|
@juliusberner Hi, I have finished the lint change. |
…ly_torch_compile All four fixtures previously called apply_torch_compile() before on_train_begin(). In the trainer (trainer.py:97,124) the order is reversed: on_train_begin() runs first (moving parameters to device/dtype and initialising preprocessors), then apply_torch_compile(). With the wrong order, preprocessors don't exist yet when compile runs and are never compiled even when torch_compile_mode is set. Also remove the now-redundant on_train_begin() calls from the two train-step tests since the fixture already covers it, and fix the ruff blank-line formatting issue.
|
@juliusberner hey Julius could you review my PR? |
2 participants
FastGen currently relies on diffusers-based model execution, which leaves performance on the table during training.
This PR adds an opt-in
torch_compile_modeflag that wraps training networks withtorch.compile, enabling PyTorch's compiler optimizations (operator fusion, memory planning, kernel autotuning) for significant speedups on common models.Benchmark
QwenImage, 20.43B params, NVIDIA H100, bfloat16, 512×512:
torch.compile(max-autotune)Speedup: 1.55x (55% faster)
Compiled iterations also show much lower variance (0.014s vs 0.094s), meaning more consistent training throughput. The one-time compilation overhead (~5–10 min with max-autotune) is amortized over the full training run.
Changes
torch_compile_mode: Optional[str] = Noneconfig option inBaseModelConfigapply_torch_compile()inFastGenModelthat compiles training modules (excluding EMA networks)apply_torch_compile()inDMD2Modelto also compileteacherandfake_scorenetworksapply_torch_compile()after DDP/FSDP wrapping inTrainer.run()apply_torch_compile()insetup_inference_modules()for the inference pathUsage
Set
torch_compile_modeto a validtorch.compilemode string (e.g.,"default"or"max-autotune") in the model config to enable. Set toNone(default) to disable.